import sys
import os
import re
import time
from PyQt5.QtWidgets import (QApplication, QMainWindow, QPushButton, QTextEdit,
                             QVBoxLayout, QHBoxLayout, QFileDialog, QWidget,
                             QLabel, QProgressBar, QStatusBar, QMessageBox,
                             QSplitter)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from openai import OpenAI


class TranslationThread(QThread):
    progress_signal = pyqtSignal(int, str)
    batch_translated_signal = pyqtSignal(str, str)  # Original, Translated
    finished_signal = pyqtSignal(str)
    error_signal = pyqtSignal(str)

    def __init__(self, file_content, prompt, max_retries=3):
        super().__init__()
        self.file_content = file_content
        self.prompt = prompt
        self.max_retries = max_retries
        self.client = OpenAI(
            api_key="0f56e8eb-8815-482d-8099-c5bef02f129a",
            base_url="https://ark.cn-beijing.volces.com/api/v3",
        )

    def split_content(self, content):
        """Split content by \section tags, only sections will be translated."""
        # 按照\section标签分割文档
        # 1. 先找出所有的section开始位置
        section_pattern = re.compile(r'\\section(\*)?(\[[^\]]*\])?{([^}]*)}', re.DOTALL)
        section_starts = list(section_pattern.finditer(content))

        if not section_starts:
            # 如果没有找到任何section，返回空列表
            return []

        batches = []
        non_section_parts = []

        # 提取每个section的内容
        for i, section_match in enumerate(section_starts):
            section_start_pos = section_match.start()
            section_title = section_match.group(3)  # section标题

            # 确定section结束位置
            if i < len(section_starts) - 1:
                section_end_pos = section_starts[i + 1].start()
            else:
                section_end_pos = len(content)

            # 提取完整section内容
            section_content = content[section_start_pos:section_end_pos]

            # 如果是第一个section，前面的内容不需要翻译
            if i == 0 and section_start_pos > 0:
                non_section_parts.append(("non_section_start", content[:section_start_pos]))

            # 添加当前section作为一个批次
            batches.append(("section", section_content, section_title))

            # 如果是最后一个section，后面的内容不需要翻译
            if i == len(section_starts) - 1 and section_end_pos < len(content):
                non_section_parts.append(("non_section_end", content[section_end_pos:]))

        # 组合section和non-section部分，按正确顺序
        all_parts = []

        # 添加文档开始的非section部分
        start_parts = [p for p in non_section_parts if p[0] == "non_section_start"]
        if start_parts:
            all_parts.extend(start_parts)

        # 添加所有section
        all_parts.extend([("section", section_content, section_title) for _, section_content, section_title in batches])

        # 添加文档结尾的非section部分
        end_parts = [p for p in non_section_parts if p[0] == "non_section_end"]
        if end_parts:
            all_parts.extend(end_parts)

        return all_parts

    def run(self):
        try:
            parts = self.split_content(self.file_content)
            # 计算需要翻译的section数量
            sections = [p for p in parts if p[0] == "section"]
            total_sections = len(sections)

            if total_sections == 0:
                self.error_signal.emit("未在文档中找到任何section标签，无法进行翻译。")
                return

            # 用于重建完整文档的字典
            translated_parts = {}
            section_count = 0

            for i, part in enumerate(parts):
                part_type = part[0]

                # 非section部分不需要翻译
                if part_type != "section":
                    content = part[1]
                    # 直接添加到翻译部分，无需翻译
                    translated_parts[i] = content
                    continue

                # 处理section部分
                section_content = part[1]
                section_title = part[2]
                section_count += 1

                retry_count = 0
                success = False

                while retry_count < self.max_retries and not success:
                    try:
                        self.progress_signal.emit(int((section_count / total_sections) * 100),
                                                  f"正在翻译第 {section_count}/{total_sections} 节: {section_title}" +
                                                  (f" (重试 {retry_count + 1})" if retry_count > 0 else ""))

                        # 准备本批次的消息
                        messages = [
                            {"role": "system", "content": self.prompt},
                            {"role": "user", "content": f"请翻译以下LaTeX内容，保留所有LaTeX格式不变，只翻译文本部分：\n{section_content}"}
                        ]

                        # 调用API
                        completion = self.client.chat.completions.create(
                            model="deepseek-r1-250120",
                            messages=messages,
                        )

                        # 获取翻译文本
                        translated_section = completion.choices[0].message.content

                        # 发送本批次翻译信号
                        self.batch_translated_signal.emit(section_content, translated_section)

                        # 添加到部分翻译结果
                        translated_parts[i] = translated_section

                        success = True

                    except Exception as e:
                        retry_count += 1
                        if retry_count >= self.max_retries:
                            self.error_signal.emit(
                                f"翻译节 {section_count}/{total_sections} 失败，已尝试 {retry_count} 次: {str(e)}")
                            # 失败时保留原文
                            translated_parts[i] = section_content
                        else:
                            # 重试前等待
                            time.sleep(2)

            # 按原始顺序重建完整文档
            full_translation = ""
            for i in range(len(parts)):
                if i in translated_parts:
                    if full_translation:
                        full_translation += translated_parts[i]
                    else:
                        full_translation = translated_parts[i]

            self.progress_signal.emit(100, "翻译完成！")
            self.finished_signal.emit(full_translation)

        except Exception as e:
            self.error_signal.emit(f"翻译过程中出错: {str(e)}")


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.file_content = ""
        self.translated_content = ""
        self.input_file_path = ""
        self.batch_translations = {}  # 存储原文和翻译对

    def init_ui(self):
        self.setWindowTitle("LaTeX文档翻译工具 - 仅翻译Section标签内容")
        self.setGeometry(100, 100, 1000, 700)

        # 主部件和布局
        main_widget = QWidget()
        main_layout = QVBoxLayout()

        # 文件选择部分
        file_section = QHBoxLayout()
        self.file_label = QLabel("未选择文件")
        self.browse_button = QPushButton("选择文件")
        self.browse_button.clicked.connect(self.browse_file)
        file_section.addWidget(self.file_label)
        file_section.addWidget(self.browse_button)

        # 翻译按钮部分
        buttons_section = QHBoxLayout()
        self.translate_button = QPushButton("开始翻译 (仅翻译\\section内容)")
        self.translate_button.clicked.connect(self.start_translation)
        self.translate_button.setEnabled(False)

        self.save_button = QPushButton("保存翻译结果")
        self.save_button.clicked.connect(self.save_translation)
        self.save_button.setEnabled(False)

        buttons_section.addWidget(self.translate_button)
        buttons_section.addWidget(self.save_button)

        # 说明标签
        description_label = QLabel("说明: 本工具将仅翻译文档中\\section标签包含的内容，保持其他部分不变")
        description_label.setStyleSheet("color: #555; font-style: italic;")

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)

        # 带分割器的文本编辑器，用于对比
        self.splitter = QSplitter(Qt.Horizontal)

        # 原文文本编辑器
        self.original_text = QTextEdit()
        self.original_text.setReadOnly(True)
        self.original_text.setPlaceholderText("原文内容将显示在这里")

        # 翻译文本编辑器
        self.translated_text = QTextEdit()
        self.translated_text.setReadOnly(True)
        self.translated_text.setPlaceholderText("翻译内容将显示在这里")

        self.splitter.addWidget(self.original_text)
        self.splitter.addWidget(self.translated_text)
        self.splitter.setSizes([500, 500])  # 初始大小相等

        # 文本编辑器标签
        editors_labels = QHBoxLayout()
        editors_labels.addWidget(QLabel("原文内容:"))
        editors_labels.addWidget(QLabel("翻译结果:"))

        # 将所有组件添加到主布局
        main_layout.addLayout(file_section)
        main_layout.addWidget(description_label)
        main_layout.addLayout(buttons_section)
        main_layout.addWidget(self.progress_bar)
        main_layout.addLayout(editors_labels)
        main_layout.addWidget(self.splitter)

        # 状态栏
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.status_bar.showMessage("准备就绪")

        # 设置主部件的布局
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

    def browse_file(self):
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(self, "选择LaTeX文件", "", "LaTeX Files (*.tex);;All Files (*)")

        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as file:
                    self.file_content = file.read()

                self.input_file_path = file_path
                self.file_label.setText(os.path.basename(file_path))
                self.translate_button.setEnabled(True)
                self.status_bar.showMessage(f"已加载文件: {os.path.basename(file_path)}")

                # 在原文文本编辑器中显示完整文件内容
                self.original_text.setText(self.file_content)
                self.translated_text.clear()  # 清除之前的翻译（如果有）
                self.batch_translations = {}  # 重置批次翻译
                self.save_button.setEnabled(False)  # 禁用保存按钮，直到新翻译完成

            except Exception as e:
                QMessageBox.critical(self, "错误", f"无法读取文件: {str(e)}")

    def start_translation(self):
        if not self.file_content:
            QMessageBox.warning(self, "警告", "请先选择一个文件。")
            return

        # 检查文件是否包含\section标签
        if not re.search(r'\\section', self.file_content):
            QMessageBox.warning(self, "警告", "未在文件中找到任何\\section标签。该工具仅翻译\\section标签内的内容。")
            return

        # 确认用户了解翻译范围
        confirm = QMessageBox.question(
            self, "确认翻译范围",
            "本工具将只翻译\\section标签内的内容，其他部分将保持原样。继续吗？",
            QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes
        )
        if confirm == QMessageBox.No:
            return

        # 翻译期间禁用按钮
        self.translate_button.setEnabled(False)
        self.browse_button.setEnabled(False)
        self.save_button.setEnabled(False)

        # 清除之前的翻译
        self.translated_content = ""
        self.translated_text.clear()
        self.batch_translations = {}

        # 创建并配置提示词，特别针对section内容翻译
        prompt = """你是一个专业的翻译助手，负责将英文LaTeX文档的section内容翻译成中文。请根据以下要求完成任务：
1. 翻译时要保留原始LaTeX文档的结构、命令、标签和其他格式元素。
2. 请仅翻译文本内容，不要改变任何LaTeX命令或环境格式（例如：\\section{Title}、\\textbf{bold text}等）。
3. LaTeX命令中的参数需要翻译，例如\\section{Introduction}应翻译为\\section{引言}。
4. 确保翻译后的中文句子与原文的语法和意思一致，并自然流畅。
5. 不要更改LaTeX代码中用于排版、引用等的任何命令。
6. 请将翻译后的内容直接返回，并保持原文的LaTeX格式不变。
7. 不要在回复中包含额外的注释或解释，只需要返回翻译后的内容。
8. 你收到的内容是一个完整的section，请确保整体翻译，不要遗漏任何部分。"""

        # 创建并启动翻译线程
        self.translation_thread = TranslationThread(self.file_content, prompt)
        self.translation_thread.progress_signal.connect(self.update_progress)
        self.translation_thread.batch_translated_signal.connect(self.batch_translated)
        self.translation_thread.finished_signal.connect(self.translation_finished)
        self.translation_thread.error_signal.connect(self.translation_error)
        self.translation_thread.start()

    def update_progress(self, value, message):
        self.progress_bar.setValue(value)
        self.status_bar.showMessage(message)

    def batch_translated(self, original, translated):
        # 存储批次翻译
        self.batch_translations[original] = translated

        # 使用所有翻译更新翻译文本编辑器
        # 需要重建完整的翻译结果
        parts = self.translation_thread.split_content(self.file_content)

        # 对于已经翻译的部分，使用翻译内容；对于尚未翻译的部分，使用原文
        current_translation = ""

        for part in parts:
            part_type = part[0]
            content = part[1]

            if part_type != "section":
                # 非section部分直接使用原文
                if current_translation:
                    current_translation += content
                else:
                    current_translation = content
            else:
                # section部分，查找是否已翻译
                if content in self.batch_translations:
                    if current_translation:
                        current_translation += self.batch_translations[content]
                    else:
                        current_translation = self.batch_translations[content]
                else:
                    # 未翻译的section使用原文
                    if current_translation:
                        current_translation += content
                    else:
                        current_translation = content

        self.translated_text.setText(current_translation)

    def translation_finished(self, translated_content):
        self.translated_content = translated_content
        self.translated_text.setText(translated_content)
        self.save_button.setEnabled(True)
        self.browse_button.setEnabled(True)
        self.translate_button.setEnabled(True)
        self.status_bar.showMessage("翻译完成！请保存结果。")

    def translation_error(self, error_message):
        QMessageBox.critical(self, "翻译错误", error_message)
        self.browse_button.setEnabled(True)
        self.translate_button.setEnabled(True)
        self.status_bar.showMessage("翻译出错，请重试。")

    def save_translation(self):
        if not self.translated_content and not self.batch_translations:
            QMessageBox.warning(self, "警告", "没有可保存的翻译内容。")
            return

        # 根据输入文件获取默认文件名
        default_filename = ""
        if self.input_file_path:
            filename, ext = os.path.splitext(self.input_file_path)
            default_filename = f"{filename}_translated{ext}"

        file_dialog = QFileDialog()
        save_path, _ = file_dialog.getSaveFileName(
            self, "保存翻译结果", default_filename, "LaTeX Files (*.tex);;All Files (*)"
        )

        if save_path:
            try:
                # 使用translated_text中的内容，它应该包含最新的翻译
                content_to_save = self.translated_text.toPlainText()

                with open(save_path, 'w', encoding='utf-8') as file:
                    file.write(content_to_save)

                self.status_bar.showMessage(f"翻译结果已保存至: {save_path}")
                QMessageBox.information(self, "成功", f"翻译结果已成功保存至:\n{save_path}")

            except Exception as e:
                QMessageBox.critical(self, "保存错误", f"无法保存文件: {str(e)}")


def main():
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()